__author__ = 'Tom Schaul, tom@idsia.ch, and Daan Wierstra'


from pybrain.rl.environments.functions.multimodal import MultiModalFunction
from scipy import sqrt, tile, swapaxes, ravel, eye, randn
import scipy 


class LennardJones(MultiModalFunction):
    """ The classical atom configuration problem. The problem dimension must be a multiple of 3, and the 
    input are the Cartesian coordinates of all atoms."""
    
    def f(self, x):
        N = self.xdim / 3
        coords = x.reshape((N, 3))
        distances = sqrt(scipy.sum((tile(coords, (N, 1, 1)) - swapaxes(tile(coords, (N, 1, 1)), 0, 1)) ** 2, axis=2)) + eye(N)
        return 2 * sum(ravel(distances ** -12 - distances ** -6))
        
    def _exampleConfig(self, numatoms, noise=0.05, edge=2.):
        """ Arranged in an approximate cube of certain edge length. """
        assert numatoms % 8 == 0
        x0 = randn(3, 2, 2, 2, numatoms / 8) * noise * edge - edge / 2
        x0[0, 0] += edge
        x0[1, :, 0] += edge
        x0[2, :, :, 0] += edge
        x0 = x0.reshape(3, numatoms).T
        return x0.flatten()
    
    @property
    def desiredValue(self):
        N = self.xdim / 3
        return self.BEST_KNOWN_TABLE[N] + 1e-5
    
    BEST_KNOWN_TABLE = {0:0, 1:0,
        2:-1.000000, 57:-288.342625,
        3:-3.000000, 58:-294.378148,
        4:-6.000000, 59:-299.738070,
        5:-9.103852, 60:-305.875476,
        6:-12.712062, 61:-312.008896,
        7:-16.505384, 62:-317.353901,
        8:-19.821489, 63:-323.489734,
        9:-24.113360, 64:-329.620147,
        10:-28.422532, 65:-334.971532,
        11:-32.765970, 66:-341.110599,
        12:-37.967600, 67:-347.252007,
        13:-44.326801, 68:-353.394542,
        14:-47.845157, 69:-359.882566,
        15:-52.322627, 70:-366.892251,
        16:-56.815742, 71:-373.349661,
        17:-61.317995, 72:-378.637253,
        18:-66.530949, 73:-384.789377,
        19:-72.659782, 74:-390.908500,
        20:-77.177043, 75:-397.492331,
        21:-81.684571, 76:-402.894866,
        22:-86.809782, 77:-409.083517,
        23:-92.844472, 78:-414.794401,
        24:-97.348815, 79:-421.810897,
        25:-102.372663, 80:-428.083564,
        26:-108.315616, 81:-434.343643,
        27:-112.873584, 82:-440.550425,
        28:-117.822402, 83:-446.924094,
        29:-123.587371, 84:-452.657214,
        30:-128.286571, 85:-459.055799,
        31:-133.586422, 86:-465.384493,
        32:-139.635524, 87:-472.098165,
        33:-144.842719, 88:-479.032630,
        34:-150.044528, 89:-486.053911,
        35:-155.756643, 90:-492.433908,
        36:-161.825363, 91:-498.811060,
        37:-167.033672, 92:-505.185309,
        38:-173.928427, 93:-510.877688,
        39:-180.033185, 94:-517.264131,
        40:-185.249839, 95:-523.640211,
        41:-190.536277, 96:-529.879146,
        42:-196.277534, 97:-536.681383,
        43:-202.364664, 98:-543.642957,
        44:-207.688728, 99:-550.666526,
        45:-213.784862, 100:-557.039820,
        46:-220.680330, 101:-563.411308,
        47:-226.012256, 102:-569.363652,
        48:-232.199529, 103:-575.766131,
        49:-239.091864, 104:-582.086642,
        50:-244.549926, 105:-588.266501,
        51:-251.253964, 106:-595.061072,
        52:-258.229991, 107:-602.007110,
        53:-265.203016, 108:-609.033011,
        54:-272.208631, 109:-615.411166,
        55:-279.248470, 110:-621.788224,
        56:-283.643105
        }




    
